from pygrib import open as pyopen
import os
import torch
import numpy as np
from utils.config import weight_root, out_path
from model.SmaAt_UNet import SmaAt_UNet
from utils.s2d import space_to_depth2
from utils.d2s import depth_to_space
# from utils.data_to_color import ImColor_tp, ImColor_tp_12mins
# import matplotlib.pyplot as plt
import glob


class reviseNet():
    def __init__(self):
        self.model = SmaAt_UNet(in_channels=4, out_channels=4)
        self.device = torch.device("cuda:0")
        weight_file = os.path.join(weight_root, f"20220304.pth")
        print(weight_file)
        checkpoints = torch.load(weight_file, map_location=self.device)
        state_dict = checkpoints['state_dict']
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

    def get_grib2(self, gfile, last=False):
        self.grbs = pyopen(gfile)
        self.grbs.seek(0)
        try:
            grb = self.grbs.select(parameterName="Total precipitation")[0]
            area_data = grb.values
            if last:
                self.startStep = grb.endStep
            else:
                self.grb = grb

            return area_data
        except ValueError:
            print("BAD GRIB2 DATA", gfile)
            return None

    def get_APCP(self, apcp_file):
        self.fdir, self.fname = os.path.split(apcp_file)

        fmins = self.fname.split("-f")[-1][:3]
        self.out_name = f"REVISED_GXR1KM_APCP_{self.fname[14:28]}-f{fmins}.grb2"

        if int(fmins) > 0:  # 排除0场数据

            last_fmin = str(int(fmins) - 12).rjust(3, '0')

            lastfile = file.replace(f"f{fmins}", f"f{last_fmin}")

            if os.path.exists(lastfile):
                try:

                    ldata = self.get_grib2(lastfile, last=True)
                    fdata = self.get_grib2(apcp_file)

                    if fdata is not None and ldata is not None:
                        apcp_data = fdata - ldata
                        apcp_data[apcp_data < 0] = 0

                        return apcp_data

                except:
                    print("Wrong", file)
                    return None

            else:
                return None
        else:
            return None

    def infer(self, afile):
        if afile[-8:-4] == 'f000':
            print(f"{os.path.basename(afile)} 0时刻数据 PASS", )
            return None

        apcp_data = self.get_APCP(afile)
        if apcp_data is None:
            print(f"BAD APCP DATA: {os.path.basename(afile)}")
            return
        revise_data = apcp_data.copy()

        valcnt = np.sum(apcp_data > 0)
        valth = apcp_data.shape[0] * apcp_data.shape[1] * 0.2

        if valcnt > valth:
            input_data = apcp_data[:880, :1160]
            x_batch = torch.from_numpy(input_data[np.newaxis, np.newaxis, ...])
            x_batch = space_to_depth2(x_batch, 2)
            with torch.no_grad():
                out = self.model(x_batch.to(device=self.device, dtype=torch.float32))
                out = depth_to_space(out, 2).cpu().numpy().squeeze()

                out[out < 0] = 0

            sub_data0 = np.maximum(out, input_data)
            sub_data1 = np.where(out < 0.6, out, sub_data0)
            revise_data[:880, :1160] = sub_data1
        else:
            print(f"{self.fname} 80% 以上区域无降水， 未订正!")

        # plt.subplot(121)
        # plt.imshow(ImColor_tp_12mins(apcp_data.copy()))
        # plt.title(f'apcp_data')
        # plt.subplot(122)
        # plt.imshow(ImColor_tp_12mins(revise_data.copy()))
        # plt.title(f'revise_data')
        # plt.gcf().set_size_inches(12, 6)
        # plt.tight_layout()
        # plt.show()
        self.grb.startStep = self.startStep

        self.grb.values = revise_data

        msg = self.grb.tostring()
        out_dir = f"{out_path}/{os.path.basename(self.fdir)}"
        if not os.path.isdir(out_dir):
            os.mkdir(out_dir)

        grbout = open(f'{out_dir}/{self.out_name}', 'wb')
        grbout.write(msg)
        grbout.close()


if __name__ == '__main__':
    # file = 'data/GX-R1KM_202106140000/Z_NAFP_C_BCGX_20210614140000_P_LARP-GXR1KM-f024.BIN'
    files = glob.glob("data/GX-R1KM_202*/Z_NAFP_C_BCGX_*.BIN")
    revise = reviseNet()
    for file in files:
        revise.infer(file)
